# Code for molecule embedding.

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Union, Tuple, Callable
from jax import nn as nn
from .embedding import Embedding
from ..model.interaction.molct_interaction import InteractionUnit
from ..utils.rbf import get_rbf, RadialBasisFunctions
from ..utils.cutoff import get_cutoff, Cutoff
from ..utils.filter import get_filter, Filter
from ..utils.activation import get_activation

class ConformationEmbedding(Embedding):

    def __init__(self,
                 fp_type = jnp.float32,
                 int_type = jnp.int32,
                 dim_node: int = 64,
                 dim_edge: int = 64,
                 num_atom_types: int = 64,
                 num_bond_types: int = 4,
                 is_emb_dis: bool = True,
                 is_emb_bond: bool = False,
                 is_interaction: bool = False,
                 dis_self: float = 0.04,
                 cutoff: float = 1.0,
                 cutoff_func: Union[str, Cutoff] = 'smooth',
                 num_basis: Optional[int] = None,
                 rbf_runc: Union[str, RadialBasisFunctions] = 'gaussian',
                 atom_filter: Union[str, Filter, None] = None,
                 bond_filter: Union[str, Filter, None] = None,
                 dis_filter: Union[str, Filter, None] = 'residual',
                 activation: Union[str, Callable] = 'silu',
                 name: str = "conformation_embedding",
                 **kwargs,
                 ):
        super().__init__(fp_type=fp_type,
                         int_type=int_type,
                         dim_node=dim_node,
                         dim_edge=dim_edge,
                         activation=activation,
                         name=name)
        
        self.is_emb_dis = is_emb_dis
        self.is_emb_bond = is_emb_bond

        # atom embedding
        self.atom_embedding = hk.Embed(vocab_size=num_atom_types,
                                       embed_dim=dim_node,
                                       w_init=hk.initializers.TruncatedNormal(1.0),
                                       name="atom_embedding") # type: ignore
        self.atom_filter = get_filter(cls_name=atom_filter,
                                      dim_in=dim_node,
                                      dim_out=dim_node,
                                      activation=self.activation,
                                      name="atom_filter")
        
        # conformation embedding
        if self.is_emb_dis:
            self.cutoff = cutoff
            self.cutoff_fn = get_cutoff(cls_name=cutoff_func,
                                        cutoff=cutoff,)
            self.rbf_fn = get_rbf(cls_name=rbf_runc,
                                  r_max=cutoff,
                                  num_basis=num_basis,)
            self.dis_filter = get_filter(cls_name=dis_filter,
                                         dim_in=self.rbf_fn.num_basis,
                                         dim_out=dim_edge,
                                         activation=self.activation,
                                         name="dis_filter")
            self.dis_self = float(dis_self)
        
        # bond embedding
        if self.is_emb_bond:
            self.bond_embedding = hk.Embed(vocab_size=num_bond_types,
                                           embed_dim=dim_edge,
                                           w_init=hk.initializers.TruncatedNormal(1.0),
                                           name="bond_embedding") # type: ignore
            self.bond_filter = get_filter(cls_name=bond_filter,
                                          dim_in=dim_edge,
                                          dim_out=dim_edge,
                                          activation=self.activation,
                                          name="bond_filter")
        
        if self.is_emb_dis and self.is_emb_bond:
            self.interaction = None
            if is_interaction:
                self.interaction = InteractionUnit(dim_feature=dim_edge,
                                                   n_heads=1,
                                                   activation=self.activation,
                                                   name="emb_interaction",) # type: ignore
        
    def __call__(self,
                 atom_type: jnp.ndarray,
                 bond_type: jnp.ndarray,
                 dist: jnp.ndarray,
                 atom_mask: Optional[jnp.ndarray] = None,
                 bond_mask: Optional[jnp.ndarray] = None,
                 dist_mask: Optional[jnp.ndarray] = None,):
        
        # atom embedding & node_embedding
        if atom_mask is None:
            atom_mask = atom_type > 0
        # (A,) -> (A, F_node)
        node_emb = self.atom_embedding(atom_type)
        # (A, F_node) -> (A, F_node)
        if self.atom_filter is not None:
            node_emb = self.atom_filter(node_emb)      
        node_mask = atom_mask

        # distance embedding
        dis_emb = None
        dis_cutoff = None
        if self.is_emb_dis:

            # set distance matrix
            # (A, A)
            eye_mask = jnp.eye(dist.shape[0], dtype=jnp.bool_)
            # (A, A)
            dist = jnp.where(eye_mask, self.dis_self, dist)

            # embedding distance
            # (A, A, B)
            dis_emb = self.rbf_fn(dist)
            if self.dis_filter is not None:
                # (A, A, B) -> (A, A, F_edge)
                dis_emb = self.dis_filter(dis_emb)
            
            # set dis_cutoff and dist_mask
            # (A, A)
            if self.cutoff_fn is None:
                dis_cutoff = jnp.ones_like(dist)
            else:
                dis_cutoff, dist_mask = self.cutoff_fn(dist, dist_mask)
        
        # distance embedding
        bond_emb = None
        bond_cutoff = None
        if self.is_emb_bond:

            # set bond type & embedding
            bond_emb = self.bond_embedding(bond_type)
            if self.bond_filter is not None:
                bond_emb = self.bond_filter(bond_emb)

            # set bond 
            if bond_mask is None:
                # (A, 1) * (1, A) -> (A, A)
                bond_mask = jnp.logical_and(jnp.expand_dims(atom_mask, axis=1), jnp.expand_dims(atom_mask, axis=0))
                # # optional bond mask
                # bond_mask = bond_type > 0

            # (A, A)
            bond_cutoff = jnp.ones_like(bond_mask)
        
        # edge embedding
        edge_emb = None
        edge_mask = None
        edge_cutoff = None
        if self.is_emb_dis and self.is_emb_bond:

            if self.interaction is not None:
                for _ in range(3):
                    node_emb, _edge_ = self.interaction(node_vec=node_emb,
                                                        edge_vec=bond_emb,
                                                        node_mask=node_mask,
                                                        edge_mask=bond_mask,
                                                        edge_cutoff=bond_cutoff,) # type: ignore
                
                edge_emb = dis_emb
                edge_cutoff = dis_cutoff
                edge_mask = dist_mask
            else:
                edge_emb = dis_emb + bond_emb
                edge_cutoff = dis_cutoff
                edge_mask = dist_mask
        
        elif self.is_emb_dis:
            edge_emb = dis_emb
            edge_mask = dist_mask
            edge_cutoff = dis_cutoff

        elif self.is_emb_bond:
            edge_emb = bond_emb
            edge_mask = bond_mask
            edge_cutoff = bond_cutoff
        else:
            raise ValueError("No edge embedding.")
        
        return node_emb, node_mask, edge_emb, edge_mask, edge_cutoff